package com.yuaer.demo.springsparkjobhandlerdemo.executor;

import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.sql.SparkSession;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

public class GlobalContext {
    // 定义SparkSession对象
    private final SparkSession spark;

    //全局参数变量
    private final Map<String, Object> globalParams;

    // 定义广播变量Map, 用于存储少量共享数据用于后续join
    private final Map<String, Broadcast<?>> broadcastVars;

    // 构造函数，传入SparkSession对象
    public GlobalContext(SparkSession spark, Map globalParams) {
        this.spark = spark;
        this.globalParams = globalParams;
        this.broadcastVars = new ConcurrentHashMap<>();
    }

    // 获取SparkSession对象
    public SparkSession getSpark() {
        return spark;
    }

    public Map<String, Object> getGlobalParams() {
        return globalParams;
    }

    // 注册广播变量，传入key和value
    public <T> void registerBroadcast(String key, T value) {
        // 使用SparkContext的broadcast方法创建广播变量
        Broadcast<T> broadcast = spark.sparkContext().broadcast(value, scala.reflect.ClassTag$.MODULE$.apply(value.getClass()));
        // 将广播变量放入Map中
        broadcastVars.put(key, broadcast);
    }
//    跟上面二选一
//    public void addBroadcast(String key, Broadcast<?> value) {
//        broadcastVars.put(key, value);
//    }

    // 获取广播变量，传入key
    @SuppressWarnings("unchecked")
    public <T> Broadcast<T> getBroadcast(String key) {
        // 从Map中获取广播变量
        return (Broadcast<T>) broadcastVars.get(key);
    }
}
